import matplotlib.pyplot as plt
import numpy as np
from dateutil import parser

from investment.time_series import IndexBasketTimeSeries, IndexTimeSeries, StockTimeSeries
from investment.operation import calc_reward_rate
from investment.date_utilities import get_date_by_year, get_date_by_month, DATE_FORMAT
from investment.constants import A_INDEX_DICT, INDEX_METRICS_FUNDAMENTAL, STOCK_CODE_DICT
from misc.plot_setting import BRIGHT_COLORS


def plot_basket(start_time: str, end_time: str):
    index_boxes = IndexBasketTimeSeries(start_time, end_time,
                                        [A_INDEX_DICT['中证500'], A_INDEX_DICT['沪深300']])
    cp = index_boxes.basket_time_series.get_fundamental_metrics('CP')
    mc = index_boxes.basket_time_series.get_fundamental_metrics('MC')
    plt.plot(index_boxes.basket_time_series.dates_series[::30], cp[::30] / cp[::30].max(), label='现价/最大')
    a_assert = index_boxes.get_merged_attribute('get_net_value')
    dyr = index_boxes.get_merged_attribute('get_dyr_increasing')
    plt.plot(index_boxes.basket_time_series.dates_series[::30], a_assert[::30] / (a_assert[::30].max()), label='资产率')
    dyr_sum = np.cumsum(dyr[::30])
    plt.plot(index_boxes.basket_time_series.dates_series[::30], dyr_sum / (dyr_sum.max()), label='股息率')
    plt.xticks(rotation=70)
    plt.legend()
    plt.show()


def plot_index_fundamental(time_series, metrics_name):
    assert metrics_name in INDEX_METRICS_FUNDAMENTAL.__members__.keys()
    assert isinstance(time_series, IndexTimeSeries)
    cp = time_series.get_fundamental_metrics('CP')
    dates = time_series.dates_series
    pe = time_series.get_fundamental_metrics(metrics_name)
    fig, ax = plt.subplots()
    ax.plot(dates, cp, color='k', alpha=0.5, label=A_INDEX_DICT.inv[time_series.stock_code])
    ax.legend(loc='upper right')
    ax_pos = ax.twinx()
    ax_pos.plot(dates, pe, color='b', alpha=0.5, label=f"{metrics_name} Ratio(RHS)")
    pe_mean = np.mean(pe)
    pe_std = np.std(pe)
    ax_pos.axhline(pe_mean, color='purple', label=f'{metrics_name} mean')
    ax_pos.axhline(pe_mean + 2 * pe_std, color='#33A1DC')
    ax_pos.axhline(pe_mean - 2 * pe_std, color='#33A1DC', label=f'{metrics_name} 2std')
    ax_pos.axhline(pe_mean + 3 * pe_std, color='#2A85B6')
    # ax_pos.axhline(pe_mean - 3 * pe_std, color='#2A85B6', label=f'{metrics_name} 3std')
    ax_pos.legend(loc='upper right', bbox_to_anchor=(1, 0.95))
    res_idx, res_date = get_date_by_year(dates)
    plt.xticks(res_idx, res_date)
    plt.xticks(rotation=70)
    plt.show()


def plot_indexes_fundamental(start_time, end_time, metrics_name):
    assert metrics_name in INDEX_METRICS_FUNDAMENTAL.__members__.keys()
    color_map = BRIGHT_COLORS
    idx = 0
    for key, value in A_INDEX_DICT.items():
        index = IndexTimeSeries(value, start_time, end_time, fill_flag=True)
        pe = index.get_fundamental_metrics(metrics_name)
        res_idx, res_date = get_date_by_year(index.dates_series)
        plt.xticks(res_idx, [date.strftime("%Y-%m") for date in res_date])
        plt.plot(index.dates_series, pe, alpha=0.6, color=color_map[idx], label=key)
        idx += 1
    plt.legend()
    plt.show()


def plot_my_stocks(start_time, end_time):
    for key, value in STOCK_CODE_DICT.items():
        code = value.code
        industry_type = value.industry_type
        stock_time_series = StockTimeSeries(code, time_start=start_time,
                                            time_end=end_time, company_type=industry_type)
        stock_time_series.fill_metrics()
        month_index, month_dates = get_date_by_month(stock_time_series.dates_series, parser.parse(end_time).day)
        pe = stock_time_series.get_fundamental_metrics("PE")
        pe_pos = stock_time_series.get_fundamental_metrics("PE_POS")
        pb = stock_time_series.get_fundamental_metrics("PB")
        pb_pos = stock_time_series.get_fundamental_metrics("PB_POS")

        fig, ax = plt.subplots()
        ax.plot(month_dates, pe[month_index], color='orange', label='pe')
        ax.plot(month_dates, pb[month_index], color='purple', label='pb')
        ax.legend(loc='upper left')
        ax_pos = ax.twinx()
        ax_pos.plot(month_dates, pe_pos[month_index], color='r', label='pe_pos')
        ax_pos.plot(month_dates, pb_pos[month_index], color='b', label='pb_pos')
        ax_pos.legend(loc='upper right')
        plt.title(value.name)
        plt.xticks(rotation=70)
        res_idx, res_date = get_date_by_year(stock_time_series.dates_series)
        plt.xticks(res_idx, [date.strftime("%Y-%m") for date in res_date])
        plt.show()


def plot_a_market_all(start_time, end_time):
    """ 验证通过 沪深A与理杏仁绘图一样， 注意是市值加权 净资产"""
    a_market = IndexTimeSeries(A_INDEX_DICT['沪深A股'], time_start=start_time,
                               time_end=end_time, fill_flag=True)
    # 需要算净值,将起点放在一起
    cp = a_market.get_fundamental_metrics("CP")
    pe = (a_market.get_fundamental_metrics("PE") + a_market.get_fundamental_metrics("PB")) / 2.
    fig, ax = plt.subplots()
    ax.plot(a_market.dates_series, cp, color='orange', label='CP')
    # ax_pos = ax.twinx()
    net_asset = a_market.get_net_asset()
    net_asset_norm = (net_asset - net_asset[0]) / (np.max(net_asset) - net_asset[0])
    asset_increasing_rate = 1.06  # asset year increasing rate to 6% year rate.
    # normalize asset to cp axis; max to increasing rate
    time_year = (parser.parse(end_time) - parser.parse(start_time)).days // 365
    cp_asset = net_asset_norm * (cp[0] * (asset_increasing_rate**time_year) - cp[0])
    cp_asset = cp_asset + cp[0]
    ax.plot(a_market.dates_series, cp_asset, color='b', label='net asset')
    ax.legend(loc='upper right')
    res_idx, res_date = get_date_by_year(a_market.dates_series)
    plt.xticks(res_idx, [date.strftime("%Y-%m") for date in res_date])

    reward_rate = calc_reward_rate(cp[0], cp[-1])
    asset_rate = calc_reward_rate(net_asset[0], net_asset[-1])
    estimation = a_market.get_thermometer()
    estimation_rate = estimation[-1] - estimation[0]
    title = f"沪深A股累计收益率:{reward_rate:.2%} , 净资产增长率:{asset_rate:.2%} , 估值变化率:{estimation_rate:.2%}"
    plt.title(title)
    plt.show()


if __name__ == "__main__":
    import datetime
    plt.rcParams['font.sans-serif'] = ['SimHei']
    today = datetime.datetime.now()
    ten_year_day = today.replace(year=today.year - 20)
    # plot_a_market_all("2011-01-01", "2021-10-01")
    # plot_basket(ten_year_day.strftime(DATE_FORMAT), today.strftime(DATE_FORMAT))
    plot_my_stocks(ten_year_day.strftime(DATE_FORMAT), today.strftime(DATE_FORMAT))
    # index_boxes = IndexTimeSeries(A_INDEX_DICT['沪深300'], "2011-10-01", "2021-10-01", fill_flag=True)
    # plot_indexes_fundamental("2011-10-01", "2021-10-01", 'PE')
    pass
